import time
from typing import List

from langchain_core.callbacks import CallbackManagerForChainRun
from langchain_core.documents import Document
from langchain_core.runnables import RunnableConfig
from loguru import logger
from starlette.responses import JSONResponse

from ai_engine.core.model import result
from ai_engine.core.model.chat import CompletionRequest
from ai_engine.core.model.vector_doc import VectorInsertOrUpdateRequest, VectorRemoveRequest
from ai_engine.core.vectorstore.client.vs_client_factory import VsClientFactory
from ai_engine.core.base.base_vs import BaseVector
from ai_engine.core.vectorstore.vector_data.knowledge.knowledge_vs_adapter import KnowledgeVectorAdapter


class KnowledgeVs(BaseVector):
    """知识库向量存储操作类"""

    def _vs_client(self):
        return VsClientFactory.get_instance(model_kwargs=self.model_kwargs, collection_name="qa")

    def insert(self, requests: List[VectorInsertOrUpdateRequest]) -> JSONResponse:
        """插入知识库定义"""
        start_time = time.time()
        # 参数验证
        KnowledgeVectorAdapter.chat_insert_validata(self.api_type, requests)
        vs_client = self._vs_client()
        insert_filter = KnowledgeVectorAdapter.insert_adapter(self.api_type, requests, vs_client.get_vector_type())
        if insert_filter is not None:
            add_doc = []
            lists = vs_client.query_documents(insert_filter.kwargs)
            if isinstance(lists, List) and len(lists) > 0:
                exist_ids = [doc.metadata["id"] for doc in lists]
                for doc in insert_filter.docs:
                    if doc.metadata["id"] not in exist_ids:
                        add_doc.append(doc)
            vs_client.add_documents(add_doc)
            end_time = time.time()
            logger.info("插入知识库，耗时：%.2f秒" % (end_time - start_time))
        return result.success(True)

    def update(self, requests: List[VectorInsertOrUpdateRequest]) -> JSONResponse:
        start_time = time.time()
        """修改知识库"""
        # 参数验证
        KnowledgeVectorAdapter.chat_insert_validata(self.api_type, requests)
        # 删除
        ids = [request.id for request in requests]
        res = self._delete(VectorRemoveRequest(ids=ids))
        if res:
            # 插入
            self.insert(requests)
            end_time = time.time()
            logger.info("更新知识库成功，耗时：%.2f秒" % (end_time - start_time))
        else:
            raise Exception("更新知识库失败，删除操作时发生异常")
        return result.success(res)

    def _delete(self, request: VectorRemoveRequest) -> bool:
        # 参数验证
        start_time = time.time()
        if len(request.ids) < 1 and len(request.categories) < 1:
            raise ValueError(f"error,id or categories must provide at least one.")
        vs_client = self._vs_client()
        delete_adapter = KnowledgeVectorAdapter.delete_adapter(request, vs_client.get_vector_type())
        res = vs_client.delete_documents(delete_adapter)
        end_time = time.time()
        logger.info("删除知识库完成，res={0}，耗时={1}".format(res, "%.2f秒" % (end_time - start_time)))
        return res

    def remove(self, request: VectorRemoveRequest) -> JSONResponse:
        """删除知识库"""
        res = self._delete(request)
        return result.success(res)

    def similarity_search(
            self,
            question: str,
            request: CompletionRequest,
            run_manager: CallbackManagerForChainRun,
    ) -> List[Document]:
        """相似度查询"""
        start_time = time.time()
        logger.info("文本生成处理->查询知识库，question=" + question)
        docs = []
        try:
            vs_client = self._vs_client()
            # 向量数据库搜索条件
            vector_search_kwargs = (KnowledgeVectorAdapter
                                    .search_kwargs_adapter(question=question,
                                                           request=request,
                                                           vector_type=vs_client.get_vector_type()))
            retriever = vs_client.get_retriever(search_kwargs=vector_search_kwargs)

            docs = retriever.invoke(
                question, config=RunnableConfig(callbacks=run_manager.get_child())
            )
            return docs
        except Exception as e:
            logger.error("文本生成处理->查询对话知识库异常，error=" + str(e), exc_info=True)
        end_time = time.time()
        logger.info("文本生成处理->查询知识库完成，耗时：%.2f秒" % (end_time - start_time))
        return docs
